﻿#ifndef RBTREE_HPP
#define RBTREE_HPP

#include <utility>
#include <cassert>
#include <iostream>
enum class Color
{
	RED,
	BLACK
};

template<class k, class value>
struct RBNode
{
	std::pair<k, value> m_kv;
	RBNode<k, value>* m_left;
	RBNode<k, value>* m_right;
	RBNode<k, value>* m_parent;
	Color m_color;

	RBNode(const std::pair<k, value>& kv)
		:m_kv(kv), m_left(nullptr), m_right(nullptr), m_parent(nullptr), m_color(Color::RED)
	{}
};

template<class k, class value>
class RBTree
{
	typedef RBNode<k, value> Node;
	typedef std::pair<k, value> PKV;
private:
	Node* m_root = nullptr;

private:
    // 父亲是祖父的左孩子, C是父亲的左孩子(需要使用右单旋)
    void rb_rotate_ll(Node* parent);
    // 父亲是祖父的右孩子, C是父亲的右孩子(需要使用左单旋)
    void rb_rotate_rr(Node* parent);


    // 检查颜色
    bool checkColor(Node* root, int count, int blackCount);
public:
    bool insert(const PKV& kv);
    bool isBalance();
};

// 父亲是祖父的左孩子，C是父亲的左孩子(需要使用右单旋)
template<class k, class value>
void RBTree<k, value>::rb_rotate_ll(Node* parent)
{
    Node* grandFather = parent->m_parent;
    Node* cur = parent->m_left;

    // 1. 孩子的右给父亲的左
    parent->m_left = cur->m_right;
    if(cur->m_right)
    {
        cur->m_right->m_parent = parent;
    }

    // 2. 孩子的右指向父亲
    cur->m_right = parent;
    parent->m_parent = cur;

    // 3. 更新cur的父子关系
    cur->m_parent = grandFather;
    if(grandFather)
    {
        if(grandFather->m_left == parent)
        {
            grandFather->m_left = cur;
        }
        else
        {
            grandFather->m_right = cur;
        }
    }
    else
    {
        m_root = cur;
    }
}

// 父亲是祖父的右孩子，C是父亲的右孩子(需要使用左单旋)
template<class k, class value>
void RBTree<k, value>::rb_rotate_rr(Node* parent)
{

    Node* grandFather = parent->m_parent;
    Node* cur = parent->m_right;

    // 1. 孩子的左给父亲的右
    parent->m_right = cur->m_left;
    if(cur->m_left)
    {
        cur->m_left->m_parent = parent;
    }

    // 2. 孩子的左指向父亲
    cur->m_left = parent;
    parent->m_parent = cur;

    // 3. 更新cur的父子关系
    cur->m_parent = grandFather;
    if(grandFather)
    {
        if(grandFather->m_right == parent)
        {
            grandFather->m_right = cur;
        }
        else
        {
            grandFather->m_left = cur;
        }
    }
    else
    {
        m_root = cur;
    }
}



template<class k, class value>
bool RBTree<k, value>::insert(const PKV& kv)
{
    // 1. 如果根节点为空，则直接插入
    if (m_root == nullptr)
    {
        m_root = new Node(kv);
        m_root->m_color = Color::BLACK;
        return true;
    }

    // 2. 如果根节点不为空，则需要找到插入位置
    Node* cur = m_root;
    Node* parent = cur->m_parent;
    while (cur)
    {
        if(kv.first > cur->m_kv.first)
        {
            parent = cur;
            cur = cur->m_right;
        }
        else if(kv.first < cur->m_kv.first)
        {
            parent = cur;
            cur = cur->m_left;
        }
        else
        {
            // 插入值已存在
            return false;
        }
    }

    // 3. 插入新节点
    Node* child = new Node(kv);
    if(kv.first > parent->m_kv.first)
    {
        parent->m_right = child;
    }
    else if(kv.first < parent->m_kv.first)
    {
        parent->m_left = child;
    }
    child->m_parent = parent;

    // 4. 调整红黑树
    // 情况1： 父亲为黑, 不用调整
    if(parent->m_color == Color::BLACK)
    {
        return true;
    }
    // 情况2： 父亲为红，需求调整(while循环向上调整)
    while(parent && parent->m_color == Color::RED)
    {
        Node* grandFather = parent->m_parent;
        // 父亲是祖父的左孩子
        if(grandFather->m_left == parent)
        {
            Node* uncle = grandFather->m_right;
            if(uncle && uncle->m_color == Color::RED)
            {
                // 情况2.1： 叔叔节点存在且为红色,则变色
                parent->m_color = Color::BLACK;
                uncle->m_color = Color::BLACK;
                grandFather->m_color = Color::RED; // 这里的grandFather可能是根节点

                // 继续向上调整
                child = grandFather;
                parent = child->m_parent;
            }
            else
            {
                // 情况2.2： 叔叔节点不存在或为黑色, 则需要旋转
                if(child == parent->m_left)
                {
                    // P是G的左孩子，C是P的左孩子，需要使用右单旋
                    //         G          P（黑色）
                    //       P  U  ==>  C  G
                    //     C                U
                    rb_rotate_ll(grandFather);
                    // 变色
                    parent->m_color = Color::BLACK;
                    child->m_color = Color::RED;

                }
                else
                {
                    // P是G的左孩子，C是P的右孩子，需要使用左右双旋
                    //         G           G         C （黑色）
                    //       P  U  ==>   C  U ==>  P  G
                    //        C         P              U
                    rb_rotate_rr(parent);
                    rb_rotate_ll(grandFather);
                    // 变色
                    parent->m_color = Color::RED;
                    grandFather->m_color = Color::RED;
                    child->m_color = Color::BLACK;

                }
            }
        }
        else
        {
            // 父亲是祖父的右孩子
            Node* uncle = grandFather->m_left;
            if(uncle && uncle->m_color == Color::RED)
            {
                // 情况2.1： 叔叔节点存在且为红色,则变色
                parent->m_color = Color::BLACK;
                uncle->m_color = Color::BLACK;
                grandFather->m_color = Color::RED; // 这里的grandFather可能是根节点

                // 继续向上调整
                child = grandFather;
                parent = child->m_parent;
            }
            else
            {
                // 情况2.2： 叔叔节点不存在或为黑色, 则需要旋转
                if(child == parent->m_right)
                {
                    // P是G的右孩子，C是P的右孩子，需要使用左单旋
                    //         G           P（黑色）
                    //       U  P  ==>   G  C
                    //           C      U  
                    rb_rotate_rr(grandFather);
                    // 变色
                    parent->m_color = Color::BLACK;
                    grandFather->m_color = Color::RED;
                    child->m_color = Color::RED;


                }
                else
                {
                    // P是G的右孩子，C是P的左孩子，需要使用右左双旋
                    //         G            G          C （黑色）
                    //       U   P  ==>   U  C  ==>  G  P
                    //          C             P     U
                    rb_rotate_rr(parent);
                    rb_rotate_ll(grandFather);
                    // 变色
                    parent->m_color = Color::RED;
                    grandFather->m_color = Color::RED;
                    child->m_color = Color::BLACK;
                }
            }
        }
    }
    // 上述调整，可能将根节点变红，因此需要将根节点变黑
    m_root->m_color = Color::BLACK;
    return true;
}


template<class k, class value>
bool RBTree<k, value>::checkColor(Node* root, int count, int blackCount)
{
    if(root == nullptr)
    {
        if(count != blackCount)
        {
            std::cout << "黑色节点数量不一致, count:" << count << " blackCount:" << blackCount << std::endl;
            return false;
        }
        return true;
    }

    // 1. 根据规则：如果节点为红，其子节点必须为黑，因此检查是否有连续的红色
    if(root->m_color == Color::RED && root->m_parent && root->m_parent->m_color == Color::RED)
    {
        std::cout << "连续的红色节点, key:" << root->m_kv.first << std::endl;
        return false;
    }
    // 2. 检查黑色节点数量
    if(root->m_color == Color::BLACK)
    {
        count++;
    }

    return checkColor(root->m_left, count, blackCount) && checkColor(root->m_right, count, blackCount);
}




template<class k, class value>
bool RBTree<k, value>::isBalance()
{
    // 1. 如果根节点为空，则直接返回true
    if(m_root == nullptr)
    {
        return true;
    }

    // 2. 检查根节点
    if(m_root->m_color != Color::BLACK)
    {
        std::cout << "根节点为红色, key:" << m_root->m_kv.first << std::endl;
        return false;
    }


    // 3. 检查颜色
    int blackCount = 0;
    Node* cur = m_root;
    while(cur)
    {
        if(cur->m_color == Color::BLACK)
        {
            blackCount++;
        }
        cur = cur->m_left;
    }

    int count = 0;
    return checkColor(m_root, count, blackCount);
}





#endif // RBTREE_HPP

